import os
import json
import base64
from pathlib import Path
from typing import Dict, List

import cv2
import numpy as np
import PIL.Image
from dashscope import MultiModalConversation

# 配置
CONFIG: Dict = {
    "model_name": "qwen-vl-max",  # 或使用 qwen-vl-plus (更快但精度稍低)
    "defect_classes": ["Corrosion", "Crack", "Peeled Paint", "Damage", "Coating Degradation"],
    "enable_preprocess": True,
    "min_confidence": 0.3,  # 降低到0.3，提高召回率
    "min_box_area": 0.00005,  # 更小的最小面积
    "max_box_area": 0.95,
    "merge_nearby_defects": True,
    "merge_distance_threshold": 0.10,
}


def preprocess_image(image: PIL.Image.Image) -> PIL.Image.Image:
    """增强图像以突出缺陷特征"""
    if not CONFIG["enable_preprocess"]:
        return image.convert("RGB")

    rgb = np.array(image.convert("RGB"))
    
    # 双边滤波：保留边缘的同时去噪
    denoised = cv2.bilateralFilter(rgb, d=5, sigmaColor=30, sigmaSpace=30)
    
    # LAB色彩空间增强
    lab = cv2.cvtColor(denoised, cv2.COLOR_RGB2LAB)
    l, a, b = cv2.split(lab)
    
    # L通道（亮度）增强
    clahe_l = cv2.createCLAHE(clipLimit=3.8, tileGridSize=(8, 8))
    l = clahe_l.apply(l)
    
    # A通道（绿-红）增强锈蚀特征
    clahe_a = cv2.createCLAHE(clipLimit=3.5, tileGridSize=(6, 6))
    a = clahe_a.apply(a)
    
    # B通道（蓝-黄）增强锈蚀特征
    clahe_b = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(6, 6))
    b = clahe_b.apply(b)
    
    enhanced = cv2.cvtColor(cv2.merge((l, a, b)), cv2.COLOR_LAB2RGB)
    
    # 锐化
    blur = cv2.GaussianBlur(enhanced, (0, 0), sigmaX=1.2)
    sharpened = cv2.addWeighted(enhanced, 1.8, blur, -0.8, 0)
    
    # 饱和度提升
    hsv = cv2.cvtColor(sharpened, cv2.COLOR_RGB2HSV)
    h, s, v = cv2.split(hsv)
    s = cv2.multiply(s, 1.15)
    s = np.clip(s, 0, 255).astype(np.uint8)
    final = cv2.cvtColor(cv2.merge((h, s, v)), cv2.COLOR_HSV2RGB)
    
    return PIL.Image.fromarray(final)


def build_prompt() -> str:
    """构建适合Qwen的简洁提示词"""
    return """分析这张风力发电叶片图片，找出所有表面缺陷并标注位置。

图片中可能存在以下缺陷（请全部找出）：

1. Corrosion（锈蚀）- 任何褐色、棕色、锈色的区域
2. Peeled Paint（涂层剥落）- 白色、浅色的露底区域
3. Crack（裂纹）- 线状裂缝
4. Damage（损伤）- 破损、凹陷
5. Coating Degradation（涂层退化）- 表面老化区域

要求：
- 仔细检查整张图片的每个部分
- 找出所有可疑的变色、破损、裂纹
- 任何褐色斑块都标记为Corrosion
- 任何白色露底都标记为Peeled Paint
- 给出每个缺陷的矩形框位置（归一化坐标0-1）

返回JSON格式：
[
  {"defect_type": "Corrosion", "bounding_box": {"x1": 0.1, "y1": 0.2, "x2": 0.3, "y2": 0.4}, "confidence": 0.8}
]

只返回JSON，不要任何解释！"""


def image_to_base64(image: PIL.Image.Image) -> str:
    """将PIL图像转换为base64字符串"""
    import io
    buffer = io.BytesIO()
    image.save(buffer, format='JPEG', quality=95)
    return base64.b64encode(buffer.getvalue()).decode('utf-8')


def call_model(image: PIL.Image.Image, api_key: str) -> List[Dict] | None:
    """调用Qwen视觉模型进行缺陷检测"""
    try:
        # 使用HTTP代理（如果需要）
        import dashscope
        dashscope.api_key = api_key
        
        # 转换为base64（避免上传到OSS）
        import io
        import base64
        buffer = io.BytesIO()
        image.save(buffer, format='JPEG', quality=95)
        image_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
        
        messages = [
            {
                "role": "user",
                "content": [
                    {"image": f"data:image/jpeg;base64,{image_base64}"},
                    {"text": build_prompt()}
                ]
            }
        ]
        
        # 先让模型描述看到了什么，激发其视觉分析能力
        response = MultiModalConversation.call(
            model=CONFIG["model_name"],
            messages=messages,
            temperature=0.01,  # 极低温度，更确定性
            top_p=0.8,
            max_tokens=2000
        )
        
        if response.status_code != 200:
            print(f"API调用失败: {response.code} - {response.message}")
            return None
        
        # 解析响应
        result_text = response.output.choices[0].message.content[0]["text"]
        
        # 保存完整调试信息
        debug_info = f"=== 完整API响应 ===\n{result_text}\n\n=== Response对象 ===\n{response}\n"
        Path("api_response_debug.txt").write_text(debug_info, encoding="utf-8")
        print(f"      调试: 原始响应长度 {len(result_text)} 字符")
        
        # 清理和解析JSON
        result_text = result_text.strip()
        if result_text.startswith("```"):
            lines = result_text.split('\n')
            result_text = '\n'.join(lines[1:-1])
        
        # 提取JSON数组
        start_idx = result_text.find('[')
        end_idx = result_text.rfind(']')
        if start_idx != -1 and end_idx != -1:
            result_text = result_text[start_idx:end_idx+1]
        
        defects = json.loads(result_text)
        print(f"      ✓ 原始检测结果: {len(defects)} 个")
        
        return defects if isinstance(defects, list) else []
        
    except json.JSONDecodeError as e:
        print(f"JSON解析失败: {e}")
        return None
    except Exception as e:
        print(f"模型调用失败: {e}")
        return None


def filter_defects(defects: List[Dict]) -> List[Dict]:
    """过滤无效的检测结果"""
    valid = []
    for d in defects:
        try:
            conf = d.get("confidence", 0)
            box = d.get("bounding_box", {})
            
            if conf < CONFIG["min_confidence"]:
                continue
            
            x1, y1 = box.get("x1", 0), box.get("y1", 0)
            x2, y2 = box.get("x2", 0), box.get("y2", 0)
            
            if not (0 <= x1 < x2 <= 1 and 0 <= y1 < y2 <= 1):
                continue
            
            area = (x2 - x1) * (y2 - y1)
            if not (CONFIG["min_box_area"] <= area <= CONFIG["max_box_area"]):
                continue
            
            valid.append(d)
        except Exception:
            continue
    
    return valid


def merge_nearby_defects(defects: List[Dict]) -> List[Dict]:
    """合并相近的同类型缺陷"""
    if not CONFIG["merge_nearby_defects"] or len(defects) <= 1:
        return defects
    
    threshold = CONFIG["merge_distance_threshold"]
    
    # 按类型分组
    by_type = {}
    for d in defects:
        dtype = d["defect_type"]
        by_type.setdefault(dtype, []).append(d)
    
    merged = []
    for dtype, group in by_type.items():
        if len(group) == 1:
            merged.extend(group)
            continue
        
        # Union-Find合并
        n = len(group)
        parent = list(range(n))
        
        def find(x):
            if parent[x] != x:
                parent[x] = find(parent[x])
            return parent[x]
        
        def union(x, y):
            parent[find(x)] = find(y)
        
        # 计算中心点距离并合并
        for i in range(n):
            b1 = group[i]["bounding_box"]
            cx1 = (b1["x1"] + b1["x2"]) / 2
            cy1 = (b1["y1"] + b1["y2"]) / 2
            
            for j in range(i+1, n):
                b2 = group[j]["bounding_box"]
                cx2 = (b2["x1"] + b2["x2"]) / 2
                cy2 = (b2["y1"] + b2["y2"]) / 2
                
                dist = ((cx1-cx2)**2 + (cy1-cy2)**2)**0.5
                if dist < threshold:
                    union(i, j)
        
        # 按组合并边界框
        groups = {}
        for i in range(n):
            root = find(i)
            groups.setdefault(root, []).append(group[i])
        
        for members in groups.values():
            if len(members) == 1:
                merged.append(members[0])
            else:
                boxes = [m["bounding_box"] for m in members]
                new_box = {
                    "x1": min(b["x1"] for b in boxes),
                    "y1": min(b["y1"] for b in boxes),
                    "x2": max(b["x2"] for b in boxes),
                    "y2": max(b["y2"] for b in boxes)
                }
                avg_conf = sum(m["confidence"] for m in members) / len(members)
                merged.append({
                    "defect_type": dtype,
                    "bounding_box": new_box,
                    "confidence": avg_conf
                })
    
    return merged


def save_yolo_annotation(defects: List[Dict], output_path: Path):
    """保存YOLO格式标注"""
    lines = []
    for d in defects:
        idx = CONFIG["defect_classes"].index(d["defect_type"])
        b = d["bounding_box"]
        xc = (b["x1"] + b["x2"]) / 2
        yc = (b["y1"] + b["y2"]) / 2
        w = b["x2"] - b["x1"]
        h = b["y2"] - b["y1"]
        lines.append(f"{idx} {xc:.6f} {yc:.6f} {w:.6f} {h:.6f}")
    
    output_path.write_text("\n".join(lines), encoding="utf-8")


def draw_boxes(image_path: Path, defects: List[Dict], output_path: Path):
    """绘制检测框并保存（支持中文路径）"""
    import cv2
    import numpy as np
    
    # 使用numpy读取，支持中文路径
    try:
        with open(image_path, 'rb') as f:
            img_data = f.read()
        img_array = np.frombuffer(img_data, dtype=np.uint8)
        img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
    except Exception as e:
        print(f"      ⚠️  无法读取图片: {e}")
        return False
    
    if img is None:
        print(f"      ⚠️  图片解码失败: {image_path}")
        return False
    
    height, width = img.shape[:2]
    colors = {
        "Corrosion": (0, 0, 255),
        "Crack": (255, 0, 0),
        "Peeled Paint": (0, 255, 255),
        "Damage": (0, 165, 255),
        "Coating Degradation": (0, 255, 0)
    }
    
    for d in defects:
        box = d["bounding_box"]
        x1 = int(box["x1"] * width)
        y1 = int(box["y1"] * height)
        x2 = int(box["x2"] * width)
        y2 = int(box["y2"] * height)
        
        color = colors.get(d["defect_type"], (255, 255, 255))
        cv2.rectangle(img, (x1, y1), (x2, y2), color, 3)
        
        label = f"{d['defect_type']} {d['confidence']:.2f}"
        font_scale = 0.7
        thickness = 2
        (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, font_scale, thickness)
        
        cv2.rectangle(img, (x1, y1-th-8), (x1+tw+4, y1), color, -1)
        cv2.putText(img, label, (x1+2, y1-4), cv2.FONT_HERSHEY_SIMPLEX, 
                   font_scale, (255, 255, 255), thickness)
    
    # 确保输出目录存在
    output_path.parent.mkdir(parents=True, exist_ok=True)
    
    # 保存图片（支持中文路径）
    try:
        _, encoded_img = cv2.imencode('.jpg', img, [cv2.IMWRITE_JPEG_QUALITY, 95])
        with open(output_path, 'wb') as f:
            f.write(encoded_img.tobytes())
        print(f"      ✅ 可视化已保存: {output_path.name}")
        return True
    except Exception as e:
        print(f"      ❌ 保存失败: {e}")
        return False


def main():
    """主函数"""
    # 获取API密钥
    api_key = os.environ.get("DASHSCOPE_API_KEY")
    if not api_key:
        print("❌ 错误: 未找到DASHSCOPE_API_KEY环境变量")
        print("\n请在终端设置:")
        print('$env:DASHSCOPE_API_KEY = "your-api-key"')
        return
    
    # 可以修改这里指定要检测的图片
    image_path = Path("风机叶片数据集/images/DJI_0004_02_06_png.rf.9a9b765846f875055aad354467ef760f.jpg")
    
    if not image_path.exists():
        print(f"❌ 图片不存在: {image_path}")
        return
    
    print("=" * 60)
    print("风力发电叶片缺陷检测系统 (Qwen版)")
    print("=" * 60)
    print()
    
    # 1. 读取图像
    print("[1/5] 正在预处理图像...")
    image = PIL.Image.open(image_path)
    print(f"      原始尺寸: {image.size}")
    
    # 2. 预处理
    processed = preprocess_image(image)
    print("      ✓ 图像预处理完成")
    print()
    
    # 3. AI分析
    print("[2/5] 正在调用Qwen AI模型...")
    defects = call_model(processed, api_key)
    
    if defects is None:
        print("      ❌ AI模型调用失败")
        return
    
    print(f"      ✓ AI分析完成，原始结果: {len(defects)} 个")
    print()
    
    # 4. 过滤
    print("[3/5] 正在过滤结果...")
    defects = filter_defects(defects)
    print(f"      ✓ 过滤后保留: {len(defects)} 个")
    print()
    
    # 5. 合并
    print("[4/5] 正在合并相近缺陷...")
    defects = merge_nearby_defects(defects)
    print(f"      ✓ 合并后最终: {len(defects)} 个")
    print()
    
    # 6. 保存结果
    print("[5/5] 正在保存结果...")
    if defects:
        # 统计
        from collections import Counter
        counts = Counter(d["defect_type"] for d in defects)
        print("      检测结果:")
        for dtype, count in sorted(counts.items()):
            print(f"      - {dtype}: {count} 个")
        
        # 保存YOLO标注
        label_path = image_path.with_suffix('.txt')
        save_yolo_annotation(defects, label_path)
        print(f"      ✓ 已保存标注: {label_path.name}")
        
        # 保存可视化
        viz_path = image_path.parent / f"{image_path.stem}_annotated.jpg"
        draw_boxes(image_path, defects, viz_path)
        print(f"      ✓ 已保存可视化: {viz_path.name}")
    else:
        print("      ℹ️  未检测到缺陷")
    
    print()
    print("=" * 60)
    print("✅ 处理完成!")
    print("=" * 60)


if __name__ == "__main__":
    main()

